#pragma once
#include "acl/acl.h"
#include "../include/custom_type.h"
#include "./data_utils.h"
#include "../include/kernel_const.h"

#include <algorithm>
#include <cstdlib>
#include <cstdio>
#include <time.h>
#include <sys/time.h>  // 用于Linux系统


#ifdef __cplusplus
extern "C" aclError rtGetC2cCtrlAddr(uint64_t* addr, uint32_t* len);
#endif


void calHValidM(
    layoutType layoutA, 
    half* h_A, 
    int64_t zeroPaddingM, 
    int64_t zeroPaddingK, 
    int64_t batchCount, 
    int64_t* h_validM
){

    struct timeval start, end;
    gettimeofday(&start, NULL);

    uint64_t strideA = zeroPaddingM * zeroPaddingK;
    for(uint32_t i = 0; i < batchCount; i++){
        half* h_A_cur = h_A + i * strideA;
        if(layoutA == 0){
            int left = 0, right = zeroPaddingM;
            int mid;
            while(right - left > 1){
                mid = left + (right - left) / 2;
                if( h_A_cur[mid * zeroPaddingK] != 0.0f ){
                    left = mid;
                }else{
                    right = mid;
                }
            }
            h_validM[i] = left + 1;
        }else{
            int left = 0, right = zeroPaddingM;
            int mid;
            while(right - left > 1){
                mid = left + (right - left) / 2;
                if( h_A_cur[mid] != 0.0f ){
                    left = mid;
                }else{
                    right = mid;
                }
            }
            h_validM[i] = left + 1;
        }
    }

    gettimeofday(&end, NULL);
    uint64_t elapsed = (end.tv_sec - start.tv_sec) * 1000000L + (end.tv_usec - start.tv_usec);
    printf("Elapsed time of calc: %ld us. \n", elapsed);

    // test
    printf("The calculated h_validM: \n");
    for(uint64_t i = 0; i < batchCount; i++ ){
        printf("%ld ", h_validM[i]);
    }
    printf("\n\n");
}

void calHValidNK(
    int64_t zeroPaddingM, 
    int64_t zeroPaddingN, 
    int64_t zeroPaddingK, 
    int64_t batchCount, 
    int64_t* h_validM, 
    int64_t* h_validN, 
    int64_t* h_validK, 
    uint8_t mode
){

    uint64_t inputValidMBytes = sizeof( int64_t ) * batchCount;
    uint64_t inputValidNBytes = sizeof( int64_t ) * batchCount; 
    uint64_t inputValidKBytes = sizeof( int64_t ) * batchCount;

    if(mode == 0){
        std::fill(h_validN, h_validN + batchCount, zeroPaddingN);
        std::fill(h_validK, h_validK + batchCount, zeroPaddingK);
    }else if(mode == 1){
        ACL_CHECK(aclrtMemcpy(h_validN, inputValidNBytes, h_validM, inputValidMBytes, ACL_MEMCPY_HOST_TO_HOST));
        std::fill(h_validK, h_validK + batchCount, zeroPaddingK);
    }else if(mode == 2){
        std::fill(h_validN, h_validN + batchCount, zeroPaddingN);
        ACL_CHECK(aclrtMemcpy(h_validK, inputValidKBytes, h_validM, inputValidMBytes, ACL_MEMCPY_HOST_TO_HOST));
    }else if(mode == 3){
        std::fill(h_validN, h_validN + batchCount, zeroPaddingN);
        std::fill(h_validK, h_validK + batchCount, zeroPaddingK);
    }   

    // test
    printf("h_validN: \n");
    for(uint64_t i = 0; i < batchCount; i++ ){
        printf("%ld ", h_validN[i]);
    }
    printf("\n\n");
    printf("h_validK: \n");
    for(uint64_t i = 0; i < batchCount; i++ ){
        printf("%ld ", h_validK[i]);
    }
    printf("\n\n");

}

void calHPointer(
    int64_t zeroPaddingM, 
    int64_t zeroPaddingN, 
    int64_t zeroPaddingK, 
    half *d_A,  
    half *d_B, 
    half *d_C, 
    int64_t batchCount, 
    half** h_A_pointer, 
    half** h_B_pointer, 
    half** h_C_pointer
){
    // 每个zeroPadding后的ABC矩阵的大小（单位是__fp元素个数）
    uint64_t strideSingleA = zeroPaddingM * zeroPaddingK;
    uint64_t strideSingleB = zeroPaddingK * zeroPaddingN;
    uint64_t strideSingleC = zeroPaddingM * zeroPaddingN;

    // 从第一个开始依次赋值
    h_A_pointer[0] = d_A;
    h_B_pointer[0] = d_B;
    h_C_pointer[0] = d_C;

    // 这里指针类型的表示范围是否有可能溢出？存疑
    for(uint64_t i = 1; i < batchCount; i++ ) {
         h_A_pointer[i] = h_A_pointer[i-1] + strideSingleA;
         h_B_pointer[i] = h_B_pointer[i-1] + strideSingleB;
         h_C_pointer[i] = h_C_pointer[i-1] + strideSingleC;
    }

    // test
    
}




void calIsPaddingAndLda(
    layoutType layoutA, 
    layoutType layoutB, 
    int64_t batchCount, 
    int64_t L1_M0, 
    int64_t L1_N0, 
    int64_t L1_K0, 
    int64_t zeroPaddingM, 
    int64_t zeroPaddingN, 
    int64_t zeroPaddingK, 
    uint8_t *h_is_A_padding, 
    uint8_t *h_is_B_padding, 
    int64_t *h_valid_padding_lda, 
    int64_t *h_valid_padding_ldb
){
    // 计算各个矩阵是否padding，若padding则计算lda和ldb
    for(uint64_t i = 0; i < batchCount; i++ ){
        int64_t AMajorSize = ( layoutA == RowMajor ? zeroPaddingK : zeroPaddingM );
        int64_t BMajorSize = ( layoutB == RowMajor ? zeroPaddingN : zeroPaddingK );
        int64_t AMajorPaddingSize = ( layoutA == RowMajor ? L1_K0 : L1_M0 );
        int64_t BMajorPaddingSize = ( layoutB == RowMajor ? L1_N0 : L1_K0 );
        if( AMajorSize % AMajorPaddingSize != 0 || AMajorSize >= 4096 ){
            h_is_A_padding[i] = 1;
            h_valid_padding_lda[i] = RoundUp<int64_t>( AMajorSize, AMajorPaddingSize );
        }else{
            h_is_A_padding[i] = 0;
            h_valid_padding_lda[i] = 0;
        }
        if( BMajorSize % BMajorPaddingSize != 0 || BMajorSize >= 4096 ){
            h_is_B_padding[i] = 1;
            h_valid_padding_ldb[i] = RoundUp<int64_t>( BMajorSize, BMajorPaddingSize );
        }else{
            h_is_B_padding[i] = 0;
            h_valid_padding_ldb[i] = 0;
        }
    }

    // test is_AB_padding valid_padding_ldab
    // printf("h_is_A_padding: \n");
    // for(uint64_t i = 0; i < batchCount; i++ ){
    //     printf("%ld ", h_is_A_padding[i]);
    // }
    // printf("\n\n");

    // printf("h_valid_padding_lda: \n");
    // for(uint64_t i = 0; i < batchCount; i++){
    //     printf("%ld ", h_valid_padding_lda[i]);
    // }
    // printf("\n\n");

    // printf("h_is_B_padding: \n");
    // for(uint64_t i = 0; i < batchCount; i++ ){
    //     printf("%ld ", h_is_B_padding[i]);
    // }
    // printf("\n\n");

    // printf("h_valid_padding_ldb: \n");
    // for(uint64_t i = 0; i < batchCount; i++){
    //     printf("%ld ", h_valid_padding_ldb[i]);
    // }
    // printf("\n\n");

    // test is_AB_padding valid_padding_ldab finished
}


void allocPointerPadding(
    int64_t batchCount, 
    int64_t L1_M0, 
    int64_t L1_N0, 
    int64_t L1_K0, 
    int64_t zeroPaddingM, 
    int64_t zeroPaddingN, 
    int64_t zeroPaddingK, 
    uint8_t *h_is_A_padding, 
    uint8_t *h_is_B_padding, 
    half **h_A_pointer_padding,  
    half **h_B_pointer_padding
){
    for(uint64_t i = 0; i < batchCount; i++ ){
        if( h_is_A_padding[i] ){
            half *d_A_pointer_padding_i = nullptr;
            uint64_t stridePaddingABytes = sizeof(half) * RoundUp<int64_t>(zeroPaddingM, L1_M0) * RoundUp<int64_t>(zeroPaddingK, L1_K0);
            ACL_CHECK(aclrtMalloc((void **)(&d_A_pointer_padding_i), stridePaddingABytes, ACL_MEM_MALLOC_HUGE_FIRST));
            h_A_pointer_padding[i] = d_A_pointer_padding_i;
        }else{
            h_A_pointer_padding[i] = nullptr;
        }
        if( h_is_B_padding[i] ){
            half *d_B_pointer_padding_i = nullptr;
            uint64_t stridePaddingBBytes = sizeof(half) * RoundUp<int64_t>(zeroPaddingN, L1_N0) * RoundUp<int64_t>(zeroPaddingK, L1_K0);
            ACL_CHECK(aclrtMalloc((void **)(&d_B_pointer_padding_i), stridePaddingBBytes, ACL_MEM_MALLOC_HUGE_FIRST));
            h_B_pointer_padding[i] = d_B_pointer_padding_i;
        }else{
            h_B_pointer_padding[i] = nullptr;
        }
    }

    // test h_AB_pointer_padding
    // printf("h_A_pointer_padding: \n");
    // for(uint64_t i = 0; i < batchCount; i++){
    //     printf("%s ", h_A_pointer_padding[i] == nullptr? "Null" : "NotNull");
    // }
    // printf("\n\n");

    // printf("h_B_pointer_padding: \n");
    // for(uint64_t i = 0; i < batchCount; i++){
    //     printf("%s ", h_B_pointer_padding[i] == nullptr? "Null" : "NotNull");
    // }
    // printf("\n\n");
    // test h_AB_pointer_padding finished
}


/**
 * @brief: 
 * @param [in] kernel_device: 要调用的算子
 * @param [in] layoutA: A矩阵排布格式
 * @param [in] layoutB: B矩阵排布格式
 * @param [in] zeroPaddingM: A、C矩阵零填充后的M维度
 * @param [in] zeroPaddingN: B、C矩阵零填充后的N维度
 * @param [in] zeroPaddingK: A、B矩阵零填充后的K维度
 * @param [in] batchCount: 批量矩阵乘的batch数
 * @param [in] d_validM: 每批矩阵乘的A、C矩阵M维度有效长度数组
 * @param [in] alpha: alpha*AB+beta*C
 * @param [in] d_A: 连续存储的零填充A矩阵首地址数组
 * @param [in] d_B: 连续存储的零填充B矩阵首地址数组
 * @param [in] beta:  alpha*AB+beta*C
 * @param [out] d_C: 每批矩阵乘的零填充B矩阵首地址数组
 * @param [in] stream: ACL同步流
 * @param [in] mode: 算子任务选择
 */
void kernel_host (
    void (*kernel_device)(
        layoutType,         // layoutA
        layoutType,         // layoutB
        int64_t,            // zeroPaddingM
        int64_t,            // zeroPaddingN
        int64_t,            // zeroPaddingK
        int64_t,            // batchCount
        __gm__ int64_t*,    // d_validM
        __gm__ int64_t*,    // d_validN
        __gm__ int64_t*,    // d_validK
        half,               // alpha
        __gm__ half**,      // d_A_pointer
        __gm__ half**,      // d_B_pointer
        half,               // beta
        __gm__ half**,      // d_C_pointer
        // 
        __gm__ uint8_t*, 
        __gm__ uint8_t*, 
        __gm__ half**, 
        __gm__ half**, 
        __gm__ int64_t* , 
        __gm__ int64_t* , 
        //
        __gm__ half**,      // d_AicAivWorkspace_Pointer
        uint64_t,           // fftsAddr
        uint8_t             // is_alpha1_beta0
    ), 
    layoutType layoutA, 
    layoutType layoutB, 
    int64_t zeroPaddingM, 
    int64_t zeroPaddingN, 
    int64_t zeroPaddingK, 
    int64_t batchCount, 
    int64_t *d_validM, 
    half alpha,  
    half *d_A,  
    half *d_B, 
    half beta, 
    half *d_C,
    aclrtStream &stream, 
    uint8_t mode
) {
    // 总任务：获取或计算算子需要的数据，包括分核策略相关参数

    // =================================================================
    // 这里还不算host侧内容，只是为了方便device侧算子计算，额外计算的一些参数
    // =================================================================

    // validMNK

    uint64_t inputValidMBytes = sizeof( int64_t ) * batchCount;
    uint64_t inputValidNBytes = sizeof( int64_t ) * batchCount; 
    uint64_t inputValidKBytes = sizeof( int64_t ) * batchCount;

    // h_validMNK[batchCount]在后面需要读取，所以从已传入device侧的数据读回来或在这里生成
    int64_t *h_validM = nullptr;
    int64_t *h_validN = nullptr;
    int64_t *h_validK = nullptr;

    ACL_CHECK(aclrtMallocHost((void **)(&h_validM), inputValidMBytes));
    ACL_CHECK(aclrtMallocHost((void **)(&h_validN), inputValidNBytes));
    ACL_CHECK(aclrtMallocHost((void **)(&h_validK), inputValidKBytes));

    // 有效长度计算 TODO
    // algorithm test
    
    struct timeval start, end;
    gettimeofday(&start, NULL);
    half* h_A; 
    uint64_t inputABytes = zeroPaddingM * zeroPaddingK * sizeof( half ) * batchCount;
    ACL_CHECK(aclrtMallocHost((void **)(&h_A), inputABytes));
    ACL_CHECK(aclrtMemcpy(h_A, inputABytes, d_A, inputABytes, ACL_MEMCPY_DEVICE_TO_HOST));
    calHValidM(
        layoutA, 
        h_A, 
        zeroPaddingM, 
        zeroPaddingK, 
        batchCount, 
        h_validM
    );

    gettimeofday(&end, NULL);
    uint64_t elapsed = (end.tv_sec - start.tv_sec) * 1000000L + (end.tv_usec - start.tv_usec);
    printf("Elapsed time of calc including memcpy: %ld us. \n", elapsed);

    ACL_CHECK(aclrtMemcpy(h_validM, inputValidMBytes, d_validM, inputValidMBytes, ACL_MEMCPY_DEVICE_TO_HOST));

    // test h_validMNK
    printf("h_validM: \n");
    for(uint64_t i = 0; i < batchCount; i++ ){
        printf("%ld ", h_validM[i]);
    }
    printf("\n\n");

    calHValidNK(
        zeroPaddingM, 
        zeroPaddingN, 
        zeroPaddingK, 
        batchCount, 
        h_validM, 
        h_validN, 
        h_validK, 
        mode
    );

    int64_t *d_validN = nullptr;
    int64_t *d_validK = nullptr;

    ACL_CHECK(aclrtMalloc((void **)(&d_validN), inputValidNBytes, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMalloc((void **)(&d_validK), inputValidKBytes, ACL_MEM_MALLOC_HUGE_FIRST));

    ACL_CHECK(aclrtMemcpy(d_validN, inputValidNBytes, h_validN, inputValidNBytes, ACL_MEMCPY_HOST_TO_DEVICE));
    ACL_CHECK(aclrtMemcpy(d_validK, inputValidKBytes, h_validK, inputValidKBytes, ACL_MEMCPY_HOST_TO_DEVICE));


    

    // d_A_pointer, d_B_pointer, d_C_pointer
    // 这里是将传入的分别由一个指针引用的连续的A、B、C矩阵，改为分别由多个指针引用
    half **h_A_pointer = nullptr;
    half **h_B_pointer = nullptr;
    half **h_C_pointer = nullptr;

    uint64_t ABCPointerBytes = batchCount * sizeof( half * );

    // aclrtMallocHost和aclrtMalloc强转类型不能是void *** 
    ACL_CHECK(aclrtMallocHost((void **)(&h_A_pointer), ABCPointerBytes)); 
    ACL_CHECK(aclrtMallocHost((void **)(&h_B_pointer), ABCPointerBytes));
    ACL_CHECK(aclrtMallocHost((void **)(&h_C_pointer), ABCPointerBytes));

    calHPointer(
        zeroPaddingM, 
        zeroPaddingN, 
        zeroPaddingK, 
        d_A, 
        d_B, 
        d_C, 
        batchCount, 
        h_A_pointer, 
        h_B_pointer, 
        h_C_pointer
    );


    // 将h_A_pointer、h_B_pointer、h_C_pointer拷贝到d_A_pointer、d_B_pointer、d_C_pointer
    half **d_A_pointer = nullptr;
    half **d_B_pointer = nullptr;
    half **d_C_pointer = nullptr;

    ACL_CHECK(aclrtMalloc((void **)(&d_A_pointer), ABCPointerBytes, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMalloc((void **)(&d_B_pointer), ABCPointerBytes, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMalloc((void **)(&d_C_pointer), ABCPointerBytes, ACL_MEM_MALLOC_HUGE_FIRST));

    ACL_CHECK(aclrtMemcpy(d_A_pointer, ABCPointerBytes, h_A_pointer, ABCPointerBytes, ACL_MEMCPY_HOST_TO_DEVICE));
    ACL_CHECK(aclrtMemcpy(d_B_pointer, ABCPointerBytes, h_B_pointer, ABCPointerBytes, ACL_MEMCPY_HOST_TO_DEVICE));
    ACL_CHECK(aclrtMemcpy(d_C_pointer, ABCPointerBytes, h_C_pointer, ABCPointerBytes, ACL_MEMCPY_HOST_TO_DEVICE));

    
    // =================================================================
    // 这里开始是关系到算子执行中所需要的参数 host侧？
    // =================================================================

    // L1_M0 L1_K0 L1_N0
    
    int64_t L1_M0 = L1M0;
    int64_t L1_N0 = L1N0;
    int64_t L1_K0 = L1K0;


    // test L1_M0N0K0
    printf("L1_M0 = %ld, L1_N0 = %ld, L1_K0 = %ld \n\n", L1_M0, L1_N0, L1_K0);
    // test L1_M0N0K0 finished

    // 放到device侧直接计算：d_padding_tasks_A, d_padding_tasks_B

    // host侧指针
    uint8_t *h_is_A_padding = nullptr;
    uint8_t *h_is_B_padding = nullptr;
    half **h_A_pointer_padding = nullptr;
    half **h_B_pointer_padding = nullptr;
    int64_t *h_valid_padding_lda = nullptr;
    int64_t *h_valid_padding_ldb = nullptr;

    // 各数组所占字节数
    uint64_t isPaddingBytes = batchCount * sizeof(uint8_t);
    uint64_t ABPointerPaddingBytes = batchCount * sizeof( half * );
    uint64_t validPaddingLdabBytes = batchCount * sizeof(int64_t);
    // uint64_t paddingTasksBytes = batchCount * sizeof(int64_t);

    // 分配host侧内存
    ACL_CHECK(aclrtMallocHost((void **)(&h_is_A_padding), isPaddingBytes));
    ACL_CHECK(aclrtMallocHost((void **)(&h_is_B_padding), isPaddingBytes));
    ACL_CHECK(aclrtMallocHost((void **)(&h_A_pointer_padding), ABPointerPaddingBytes));
    ACL_CHECK(aclrtMallocHost((void **)(&h_B_pointer_padding), ABPointerPaddingBytes));
    ACL_CHECK(aclrtMallocHost((void **)(&h_valid_padding_lda), validPaddingLdabBytes));
    ACL_CHECK(aclrtMallocHost((void **)(&h_valid_padding_ldb), validPaddingLdabBytes));

    // is_padding and valid_padding_ld
    calIsPaddingAndLda(
        layoutA, 
        layoutB, 
        batchCount, 
        L1_M0, 
        L1_N0, 
        L1_K0, 
        zeroPaddingM, 
        zeroPaddingN, 
        zeroPaddingK, 
        h_is_A_padding, 
        h_is_B_padding, 
        h_valid_padding_lda, 
        h_valid_padding_ldb
    );

    uint8_t *d_is_A_padding = nullptr;
    uint8_t *d_is_B_padding = nullptr;
    int64_t *d_valid_padding_lda = nullptr;
    int64_t *d_valid_padding_ldb = nullptr;

    ACL_CHECK(aclrtMalloc((void **)(&d_is_A_padding), isPaddingBytes, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMalloc((void **)(&d_is_B_padding), isPaddingBytes, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMalloc((void **)(&d_valid_padding_lda), validPaddingLdabBytes, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMalloc((void **)(&d_valid_padding_ldb), validPaddingLdabBytes, ACL_MEM_MALLOC_HUGE_FIRST));

    ACL_CHECK(aclrtMemcpy(d_is_A_padding, isPaddingBytes, h_is_A_padding, isPaddingBytes, ACL_MEMCPY_HOST_TO_DEVICE));
    ACL_CHECK(aclrtMemcpy(d_is_B_padding, isPaddingBytes, h_is_B_padding, isPaddingBytes, ACL_MEMCPY_HOST_TO_DEVICE));
    ACL_CHECK(aclrtMemcpy(d_valid_padding_lda, validPaddingLdabBytes, h_valid_padding_lda, validPaddingLdabBytes, ACL_MEMCPY_HOST_TO_DEVICE));
    ACL_CHECK(aclrtMemcpy(d_valid_padding_ldb, validPaddingLdabBytes, h_valid_padding_ldb, validPaddingLdabBytes, ACL_MEMCPY_HOST_TO_DEVICE));

    


    // pointer_padding
    allocPointerPadding(
        batchCount, 
        L1_M0, 
        L1_N0, 
        L1_K0, 
        zeroPaddingM, 
        zeroPaddingN, 
        zeroPaddingK, 
        h_is_A_padding, 
        h_is_B_padding, 
        h_A_pointer_padding, 
        h_B_pointer_padding
    );

    half **d_A_pointer_padding = nullptr;
    half **d_B_pointer_padding = nullptr;

    ACL_CHECK(aclrtMalloc((void **)(&d_A_pointer_padding), ABPointerPaddingBytes, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMalloc((void **)(&d_B_pointer_padding), ABPointerPaddingBytes, ACL_MEM_MALLOC_HUGE_FIRST));

    ACL_CHECK(aclrtMemcpy(d_A_pointer_padding, ABPointerPaddingBytes, h_A_pointer_padding, ABPointerPaddingBytes, ACL_MEMCPY_HOST_TO_DEVICE));
    ACL_CHECK(aclrtMemcpy(d_B_pointer_padding, ABPointerPaddingBytes, h_B_pointer_padding, ABPointerPaddingBytes, ACL_MEMCPY_HOST_TO_DEVICE));


    // cBlockSum, d_C_blocks
    // 算结果矩阵C按[L1_M0, L1_N0, L1_K0]来划分可分多少[L1_M0, L1_N0]任务块（注意边界处可能不足[L1_M0, L1_N0, L1_K0]）
    // 以及计算总的任务数

    int64_t cBlocksSum = 0;
    int64_t *h_C_blocks = nullptr;

    uint64_t cBlocksBytes = batchCount * sizeof(int64_t);

    ACL_CHECK(aclrtMallocHost((void **)(&h_C_blocks), cBlocksBytes));

    for(uint64_t i = 0; i < batchCount; i++ ){
        h_C_blocks[i] = CeilDiv<int64_t>(h_validM[i], L1_M0) * CeilDiv<int64_t>(h_validN[i], L1_N0);
        cBlocksSum += h_C_blocks[i];
    }

    // test h_C_blocks and cBlocksSum
    printf("h_C_blocks: \n");
    for(uint64_t i = 0; i < batchCount; i++){
        printf("%ld ", h_C_blocks[i]);
    }
    printf("\ncBlocksSum = %ld\n\n", cBlocksSum);
    // test h_C_blocks and cBlocksSum finished


    // workspace结果块双缓冲方案
    half* d_AicAivWorkspace = nullptr;
    // 这里原来是uint32_t，在部分样例下会爆32位整数范围 如MNKC=8972 8107 845 47
    // uint64_t workspaceBytes = zeroPaddingM * zeroPaddingN * sizeof( half ) * batchCount;
    uint64_t WORKSPACE_NUM = WORKSPACENUM;
    uint64_t workspaceSingleCoreSize = WORKSPACE_NUM * L1_M0 * L1_N0;
    uint64_t workspaceBytes = MAX_AICORE_NUM * workspaceSingleCoreSize * sizeof(half);
    ACL_CHECK(aclrtMalloc((void**)(&d_AicAivWorkspace), workspaceBytes, ACL_MEM_MALLOC_HUGE_FIRST));

    half** h_AicAivWorkspace_Pointer = nullptr; 
    // uint64_t workspacePointerBytes = ABCPointerBytes; // 这个方案，数组长度为矩阵数
    uint64_t workspacePointerBytes = MAX_AICORE_NUM * sizeof(half *); // 当前方案，数组长度为aicore数
    ACL_CHECK(aclrtMallocHost((void **)(&h_AicAivWorkspace_Pointer), workspacePointerBytes));

    h_AicAivWorkspace_Pointer[0] = d_AicAivWorkspace;
    for(uint64_t i = 1; i < /*batchCount*/ MAX_AICORE_NUM; i++){
        // h_AicAivWorkspace_Pointer[i] = h_AicAivWorkspace_Pointer[i-1] + strideSingleC;
        h_AicAivWorkspace_Pointer[i] = h_AicAivWorkspace_Pointer[i-1] + workspaceSingleCoreSize;
    }

    half** d_AicAivWorkspace_Pointer = nullptr;
    ACL_CHECK(aclrtMalloc((void **)(&d_AicAivWorkspace_Pointer), workspacePointerBytes, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMemcpy(d_AicAivWorkspace_Pointer, workspacePointerBytes, h_AicAivWorkspace_Pointer, workspacePointerBytes, ACL_MEMCPY_HOST_TO_DEVICE));


    // 设置AIC AIV同步需要的变量，这个地址需要传到算子，在AIC和AIV运行前设置该地址
    uint64_t fftsAddr{0};
    uint32_t fftsLen{0};
    ACL_CHECK(rtGetC2cCtrlAddr(&fftsAddr, &fftsLen));

    // 在host侧做判断（device侧无法做这个判断）
    uint8_t is_alpha1_beta0 = (alpha == 1.0f && beta == 0.0f );

    struct timeval startKernel, endKernel;
    gettimeofday(&startKernel, NULL);
    // kernel调用
    (*kernel_device)<<< MAX_AICORE_NUM, nullptr, stream >>>(
        layoutA, 
        layoutB, 
        zeroPaddingM, 
        zeroPaddingN, 
        zeroPaddingK, 
        batchCount, 
        d_validM, 
        d_validN, 
        d_validK, 
        alpha, 
        d_A_pointer, 
        d_B_pointer, 
        beta, 
        d_C_pointer, 
        // padding
        d_is_A_padding, 
        d_is_B_padding, 
        d_A_pointer_padding, 
        d_B_pointer_padding, 
        d_valid_padding_lda, 
        d_valid_padding_ldb, 
        //
        d_AicAivWorkspace_Pointer, 
        fftsAddr, 
        is_alpha1_beta0
    );

    // 同步后才能测数据
    ACL_CHECK(aclrtSynchronizeStream(stream));
    gettimeofday(&endKernel, NULL);
    uint64_t elapsedKernel = (endKernel.tv_sec - startKernel.tv_sec) * 1000000L + (endKernel.tv_usec - startKernel.tv_usec);
    printf("Elapsed time of kernel: %ld us. \n", elapsedKernel);

    ACL_CHECK(aclrtFreeHost(h_A));
    ACL_CHECK(aclrtFreeHost(h_validM));
    ACL_CHECK(aclrtFreeHost(h_validN));
    ACL_CHECK(aclrtFreeHost(h_validK));
    ACL_CHECK(aclrtFreeHost(h_A_pointer));
    ACL_CHECK(aclrtFreeHost(h_B_pointer));
    ACL_CHECK(aclrtFreeHost(h_C_pointer));

    for(uint32_t i = 0; i < batchCount; i++){
        if( h_is_A_padding[i] ){
            ACL_CHECK(aclrtFree(h_A_pointer_padding[i]));
        }
        if( h_is_B_padding[i] ){
            ACL_CHECK(aclrtFree(h_B_pointer_padding[i]));
        }
    }
    ACL_CHECK(aclrtFreeHost(h_is_A_padding));
    ACL_CHECK(aclrtFreeHost(h_is_B_padding));
    ACL_CHECK(aclrtFreeHost(h_A_pointer_padding));
    ACL_CHECK(aclrtFreeHost(h_B_pointer_padding));
    ACL_CHECK(aclrtFreeHost(h_valid_padding_lda));
    ACL_CHECK(aclrtFreeHost(h_valid_padding_ldb));

    ACL_CHECK(aclrtFreeHost(h_C_blocks));
    ACL_CHECK(aclrtFreeHost(h_AicAivWorkspace_Pointer));


    ACL_CHECK(aclrtFree(d_validN));
    ACL_CHECK(aclrtFree(d_validK));
    ACL_CHECK(aclrtFree(d_A_pointer));
    ACL_CHECK(aclrtFree(d_B_pointer));
    ACL_CHECK(aclrtFree(d_C_pointer));

    ACL_CHECK(aclrtFree(d_is_A_padding));
    ACL_CHECK(aclrtFree(d_is_B_padding));
    ACL_CHECK(aclrtFree(d_A_pointer_padding));
    ACL_CHECK(aclrtFree(d_B_pointer_padding));
    ACL_CHECK(aclrtFree(d_valid_padding_lda));
    ACL_CHECK(aclrtFree(d_valid_padding_ldb));

    // ACL_CHECK(aclrtFree(d_C_blocks));
    ACL_CHECK(aclrtFree(d_AicAivWorkspace));
    ACL_CHECK(aclrtFree(d_AicAivWorkspace_Pointer));


}